package com.aiserver.interceptor;

import com.aiserver.util.EncryptUtil;
import com.aiserver.wrapper.RepeatableReadHttpServletRequest;
import com.alibaba.fastjson.JSON;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.HashMap;
import java.util.Map;

/**
 * 加密拦截器
 * 处理请求和响应的加密解密
 */
@Component
public class EncryptionInterceptor implements HandlerInterceptor {

    private static final Logger logger = LoggerFactory.getLogger(EncryptionInterceptor.class);

    @Autowired
    private EncryptUtil encryptUtil;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        // 只处理POST和PUT请求
        String method = request.getMethod();
        if (!"POST".equals(method) && !"PUT".equals(method)) {
            return true;
        }

        // 检查是否需要解密
        String encrypted = request.getHeader("X-Encrypted");
        if (!"true".equals(encrypted)) {
            return true;
        }

        try {
            // 如果请求已经被包装过，直接使用
            String requestBody;
            if (request instanceof RepeatableReadHttpServletRequest) {
                requestBody = ((RepeatableReadHttpServletRequest) request).getBody();
            } else {
                // 这种情况不应该发生，因为Filter应该已经包装了请求
                logger.warn("请求未被正确包装，无法处理加密数据");
                return true;
            }

            if (!StringUtils.hasText(requestBody)) {
                return true;
            }

            // 解析加密数据
            Map<String, Object> encryptedData = JSON.parseObject(requestBody, Map.class);
            String encryptedContent = (String) encryptedData.get("data");
            
            if (!StringUtils.hasText(encryptedContent)) {
                logger.warn("加密请求中缺少data字段");
                response.setStatus(HttpServletResponse.SC_BAD_REQUEST);
                response.setContentType("application/json;charset=UTF-8");
                response.getWriter().write("{\"code\":400,\"message\":\"加密数据格式错误\",\"timestamp\":" + System.currentTimeMillis() + "}");
                return false;
            }

            // 解密数据
            String decryptedContent = encryptUtil.aesDecrypt(encryptedContent);
            logger.debug("请求数据解密成功");

            // 将解密后的数据存储到请求属性中
            request.setAttribute("decryptedData", decryptedContent);
            
            return true;
        } catch (Exception e) {
            logger.error("请求数据解密失败", e);
            response.setStatus(HttpServletResponse.SC_BAD_REQUEST);
            response.setContentType("application/json;charset=UTF-8");
            response.getWriter().write("{\"code\":400,\"message\":\"数据解密失败\",\"timestamp\":" + System.currentTimeMillis() + "}");
            return false;
        }
    }

    /**
     * 加密响应数据
     * 注意：这个方法需要配合ResponseBodyAdvice使用
     */
    public String encryptResponse(String responseData) {
        try {
            if (!StringUtils.hasText(responseData)) {
                return responseData;
            }
            
            String encryptedData = encryptUtil.aesEncrypt(responseData);
            
            // 构造加密响应格式
            Map<String, Object> result = new HashMap<>();
            result.put("encrypted", true);
            result.put("data", encryptedData);
            result.put("timestamp", System.currentTimeMillis());
            
            return JSON.toJSONString(result);
        } catch (Exception e) {
            logger.error("响应数据加密失败", e);
            return responseData;
        }
    }
}